//
// Created by mio on 25-5-13.
//

#include "AACEncoder.h"
#include "AudioResampler.h"

AACEncoder::AACEncoder(): audio_resampler_(nullptr) {
}

AACEncoder::~AACEncoder()
{
    Close();
}

bool AACEncoder::Open(AVConfig &config)
{
    if (is_initialzed_) return false;

    config_ = config;
    codec_ = (AVCodec*)(avcodec_find_encoder(AV_CODEC_ID_AAC));
    if (!codec_) {
        Close();
        return false;
    }

    // 创建编码器上下文
    codecContext_ = avcodec_alloc_context3(codec_);
    if (!codecContext_) {
        Close();
        return false;
    }

    // 设置编码器上下文
    codecContext_->sample_rate = config.audio.samplerate;
    codecContext_->sample_fmt = AV_SAMPLE_FMT_FLTP;
    codecContext_->channels = config.audio.channels;
    codecContext_->channel_layout = av_get_default_channel_layout(config.audio.channels);
    codecContext_->bit_rate = config.audio.bitrate;

    // 获取全局AAC头部
    codecContext_->flags |= AV_CODEC_FLAG_GLOBAL_HEADER;

    if (avcodec_open2(codecContext_, codec_, nullptr) != 0) {
        Close();
        return false;
    }

    audio_resampler_.reset(new AudioResampler());

    // 初始化该重采样对象, 将格式转换为FLTP
    if (!audio_resampler_->Open(config.audio.samplerate, config.audio.channels, config.audio.format,
        config.audio.samplerate, config.audio.channels, AV_SAMPLE_FMT_FLTP)) {
        Close();
        return false;
    }

    is_initialzed_ = true;
    return true;
}

void AACEncoder::Close() {
    if (audio_resampler_) {
        audio_resampler_->Close();
        audio_resampler_.reset();
        audio_resampler_ = nullptr;
    }
}

// TODO: 搞清楚这里的frame_size的含义
uint32_t AACEncoder::GetFrameSamples() {
    if (is_initialzed_) {
        return codecContext_->frame_size;
    }

    return 0;
}

AVPacketPtr AACEncoder::Encode(const uint8_t *pcm, int samples)
{

    // 初始化Frame
    AVFramePtr in_frame(av_frame_alloc(), [](AVFrame* ptr){ av_frame_free(&ptr); });
    in_frame->sample_rate = codecContext_->sample_rate;
    in_frame->channels = codecContext_->channels;
    in_frame->nb_samples = samples;
    in_frame->channel_layout = codecContext_->channel_layout;
    in_frame->format = AV_SAMPLE_FMT_FLTP;
    in_frame->pts = pts_;
    in_frame->pts = av_rescale_q(pts_, { 1, codecContext_->sample_rate }, codecContext_->time_base);

    // 创建内存
    if (av_frame_get_buffer(in_frame.get(), 0) < 0) return nullptr;

    // 计算位数
    int bytes_per_sample = av_get_bytes_per_sample(config_.audio.format);
    if (bytes_per_sample == 0) return nullptr;

    memcpy(in_frame->data[0], pcm, bytes_per_sample * in_frame->channels * samples);

    // 重采样
    AVFramePtr fltp_frame = nullptr;
    if(audio_resampler_->Convert(in_frame, fltp_frame) <= 0) return nullptr;

    int ret = avcodec_send_frame(codecContext_, fltp_frame.get());
    if (ret != 0) return nullptr;

    AVPacketPtr av_packet(av_packet_alloc(), [](AVPacket* p){ av_packet_free(&p); });
    av_init_packet(av_packet.get());
    ret = avcodec_receive_packet(codecContext_, av_packet.get());
    if (ret == AVERROR(EAGAIN) || ret == AVERROR_EOF) return nullptr;
    else if (ret < 0) return nullptr;

    return av_packet;
}
